package com.gitee.feizns.explore.data_structure.tree.cus.avl;

import com.gitee.feizns.explore.data_structure.tree.cus.Node;
import com.gitee.feizns.explore.data_structure.tree.cus.PutRet;
import com.gitee.feizns.explore.data_structure.tree.cus.TreeUtils;

import java.util.Collection;
import java.util.Set;
import java.util.TreeSet;
import java.util.stream.Collectors;

/**
 * @author feizns
 * @since 2019/11/28
 */
public class AvlImpl2<K extends Comparable<K>, V> implements Avl<AvlImpl2.AvlNode, K, V> {

    private AvlNode<K, V> root;

    private int size;

    @Override
    public Node<AvlNode, K, V> root() {
        return root;
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public V get(Object key) {
        AvlNode node = TreeUtils.find(root, (K) key);
        return node != null ? (V) node.val() : null;
    }

    @Override
    public V put(K key, V value) {
        PutRet<AvlNode, V> ret = put(root(), key, value);
        root = ret.node();
        return ret.v();
    }

    private PutRet<AvlNode, V> put(Node<AvlNode, K, V> node, K k, V v) {
        if ( node == null ) {
            size++;
            return new PutRet<>(new AvlNode(k, v), null);
        }

        V oldV;
        int com = k.compareTo(node.key());
        if ( com < 0 ) {
            PutRet<AvlNode, V> ret = put(node.left(), k, v);
            node.left(ret.node());
            oldV = ret.v();
        } else if ( com > 0 ) {
            PutRet<AvlNode, V> ret = put(node.right(), k, v);
            node.right(ret.node());
            oldV = ret.v();
        } else {
            oldV = node.val();
            node.val(v);
        }

        AvlNode avlNode = node.node();
        resetHeight(avlNode);

        return new PutRet<>(rotate(avlNode), oldV);
    }

    /**
     * 右旋
     * @param node
     * @return
     */
    private AvlNode lll(AvlNode node) {
        AvlNode left = node.left();
        node.left(left.right());
        left.right(node);
        resetHeight(node);
        resetHeight(left);
        return left;
    }

    /**
     * 左旋
     * @param node
     * @return
     */
    private AvlNode rrr(AvlNode node) {
        AvlNode right = node.right();
        node.right(right.left());
        right.left(node);
        resetHeight(node);
        resetHeight(right);
        return right;
    }

    private AvlNode llr(AvlNode node) {
        node.left(rrr(node.left()));
        return lll(node);
    }

    private AvlNode rrl(AvlNode node) {
        node.right(lll(node.right()));
        return rrr(node);
    }

    @Override
    public V remove(Object key) {
        PutRet<AvlNode, V> ret = remove(root, (K) key);
        root = ret.node();
        return ret.v();
    }

    private PutRet<AvlNode, V> remove(AvlNode node, K k) {
        if ( node == null ) {
            return new PutRet<>(null, null);
        }

        AvlNode avlNode = node;
        V oldV;
        int com = k.compareTo((K) node.key());
        if ( com < 0 ) {
            PutRet<AvlNode, V> ret = remove(avlNode.left(), k);
            node.left(ret.node());
            oldV = ret.v();
        } else if ( com > 0 ){
            PutRet<AvlNode, V> ret = remove(avlNode.right(), k);
            node.right(ret.node());
            oldV = ret.v();
        } else {
            oldV = (V) node.val();
            size--;
            if ( node.left() != null && node.right() != null ) {
                Node<AvlNode, K, V> min = min(node.right());
                min.right(remove(node.right(), min.key()).node());
                min.left(node.left());
                avlNode = (AvlNode) min;
            } else if ( node.left() == null ) {
                avlNode = node.right();
            } else {
                avlNode = node.left();
            }
        }

        if ( avlNode != null ) {
            resetHeight(avlNode);
            avlNode = rotate(avlNode);
        }

        return new PutRet<>(avlNode, oldV);
    }

    private AvlNode rotate(AvlNode node) {
        int dif = balancedFactor(node);
        if ( dif > 1 && balancedFactor(node.left()) < 0 )
            node = llr(node);
        else if ( dif < -1 && balancedFactor(node.right()) > 0 )
            node = rrl(node);
        else if ( dif > 1 )
            node = lll(node);
        else if ( dif < -1 )
            node = rrr(node);
        return node;
    }

    @Override
    public void clear() {
        size = 0;
        root = null;
    }

    @Override
    public Set<K> keySet() {
        return entrySet().stream().map(Entry::getKey).collect(Collectors.toSet());
    }

    @Override
    public Collection<V> values() {
        return entrySet().stream().map(Entry::getValue).collect(Collectors.toList());
    }

    @Override
    public Set<Entry<K, V>> entrySet() {
        TreeSet<Entry<K, V>> entries = new TreeSet<>();
        middleEach((k, v) -> entries.add(new EntryImpl<>(k, v)));
        return entries;
    }

    private void resetHeight(AvlNode node) {
        node.height = Math.max(height(node.left()), height(node.right())) + 1;
    }

    private int balancedFactor(AvlNode node) {
        return height(node.left()) - height(node.right());
    }

    private int height(AvlNode node) {
        return node != null ? node.height : 0;
    }

    private Node<AvlNode, K, V> min(AvlNode node) {
        return node == null || node.left() == null ? node : min(node.left());
    }

    public static class AvlNode<K extends Comparable<K>, V> implements Node<AvlNode, K, V> {

        K k;

        V v;

        public AvlNode(K k, V v) {
            this.k = k;
            this.v = v;
            this.height = 1;
        }

        int height;

        AvlNode left, right;

        @Override
        public K key() {
            return k;
        }

        @Override
        public void key(K k) {
            this.k = k;
        }

        @Override
        public V val() {
            return v;
        }

        @Override
        public void val(V v) {
            this.v = v;
        }

        @Override
        public AvlNode left() {
            return left;
        }

        @Override
        public void left(AvlNode left) {
            this.left = left;
        }

        @Override
        public AvlNode right() {
            return right;
        }

        @Override
        public void right(AvlNode right) {
            this.right = right;
        }

    }

    public static class EntryImpl<K, V> implements Entry<K, V> {

        K k;

        V v;

        public EntryImpl(K k, V v) {
            this.k = k;
            this.v = v;
        }

        @Override
        public K getKey() {
            return k;
        }

        @Override
        public V getValue() {
            return v;
        }

        @Override
        public V setValue(V value) {
            V old = v;
            v = value;
            return old;
        }

    }

}
